#!/usr/bin/python
# Prints COG statistics and also counts of individual COG category

import sys
import re

cog_cat_file = sys.argv[1]
ccf = open(cog_cat_file, "rU")
ccl = ccf.readlines()

org_list_file = sys.argv[2]
olf = open(org_list_file, "rU")
oll = olf.readlines()

genome_cog_file = sys.argv[3]
gcf = open(genome_cog_file, "rU")
gcl = gcf.readlines()

#outfile = sys.argv[4]
#of = open(outfile, "w")

cog_cats = []
cog_ids = []
cog_defs = []

for line in ccl:
    l = line.split('\t')
    cog_cat = l[0]
    cog_id = l[1]
    cog_def = l[2].rstrip()
    cog_cats.append(cog_cat)
    cog_ids.append(cog_id)
    cog_defs.append(cog_def)

orgs = {}
for line in oll:
    l = line.split('\t')
    acc = l[0].rstrip()
    name = l[1].rstrip()
    orgs[acc] = name

cogs_found = []

for line in gcl:
    l = line.split('\t')
    locus_tag = l[0]
    cog_id = l[1]
    #evalue = float(l[2].rstrip())
    cogs_found.append(cog_id)

m1 = re.compile('(.*).cogs.list')
acc_id = m1.match(genome_cog_file).group(1)
org_name = ""
if acc_id in orgs:
    org_name = orgs[acc_id]

cog_counts = []

print org_name

for x in cog_ids:
    c = cogs_found.count(x)
    cog_counts.append(str(c))
    print c

#of.write(cog_counts)

ccf.close()
olf.close()
gcf.close()
